# -*- coding:utf-8 -*-

import numpy as np
import math
import cv2


def bgr_to_yuvi420(rgb):
    r""" BGR转YUVI420 """
    return cv2.cvtColor(rgb, cv2.COLOR_BGR2YUV_I420)


def yuvi420_to_nv21(yuvi420):
    r""" YUVI420转NV21 """
    h = yuvi420.shape[0] * 2//3
    w = yuvi420.shape[1]
    Y = yuvi420[0:h]
    U = yuvi420[h:5*h//4].reshape([-1])
    V = yuvi420[5*h//4:].reshape([-1])
    VU = np.array(list(zip(V, U))).reshape([h//2, w])
    nv21 = np.concatenate((Y, VU),axis=0)
    return nv21


def read_file_yile(file_path, one_patch_size):
    r"""以one_patch_size逐个读取数据, yield为python3内容,需要for处理 """
    with open(file_path, "rb") as f:
        while True:
            patch = f.read(one_patch_size)
            if patch:
                yield patch
            else:
                return ""


def load_nv21_image(patch, w, h):
    r""" 读取NV21 """
    arr = np.frombuffer(patch, np.uint8)
    nv21 = arr.reshape([h*3//2, w])
    return nv21


def nv21_to_gray(nv21):
    r""" 取NV21的Y分量作为灰度图 """
    h = nv21.shape[0] * 2 // 3
    return nv21[0:h]


def nv21_to_bgr(nv21):
    r"""" NV21转BGR """
    return cv2.cvtColor(nv21, cv2.COLOR_YUV2BGR_NV21)


def bgr_to_rgb(bgr):
    r""" R,G通道互换 """
    return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)


r""" R,G通道互换 """
rgb_to_bgr = bgr_to_rgb   
   

def cal_mean_gradient_one_channel(im):
    r""" 统计单通道图平均梯度。仅适合灰度图或者单个分量 """
    # 计算要转到浮点去,否则uint8不具备计算梯度的能力
    Ab = im[:-1, :-1].astype(float)
    Ay = im[1:, :-1].astype(float)
    Ax = im[:-1, 1:].astype(float)
    # 梯度计算
    gy = Ab - Ay
    gx = Ab - Ax
    # 在numpy中 * 是element wide [ELT] 相乘, 并不是行列式[GEMM]
    # 梯度 = 水平梯度 和 竖直梯度 的 欧式距离
    grad =  np.sqrt((gx*gx  + gy * gy)/2)
    return np.mean(grad)


def mean_gradient(im):
    r""" 统计图像的平均梯度 """
    dim = len(im.shape)
    if dim < 2 or dim > 3:
        raise "image is offen 2 or 3 dim, here is invalid"
    ans = []
    if dim == 2:
        v = cal_mean_gradient_one_channel(im)
        ans.append(v)
    else:
        for i in range(im.shape[2]):
            v = cal_mean_gradient_one_channel(im[:, :, i])
            ans.append(v)
    return ans


def spatial_freq_one_channel(im):
    r""" 统计统计单通道图空间频率 仅适合灰度图或者单个分量 """
    # 计算要转到浮点去,否则uint8不具备计算梯度的能力
    Ab = im[:-1, :-1].astype(float)
    Ay = im[1:, :-1].astype(float)
    Ax = im[:-1, 1:].astype(float)
    # 梯度计算
    gy = Ab - Ay
    gx = Ab - Ax
    # RA平方计算,向量变标量
    RAd = np.sum(gy * gy)/(im.shape[0] * im.shape[1])
    # RA平方计算,向量变标量
    CAd = np.sum(gx * gx)/(im.shape[0] * im.shape[1])
    # sf计算,这里np.sqrt也可以计算标量,用math的是为了便于理解
    SF = math.sqrt(RAd +CAd)
    return SF
   
 
def spatial_freq(im):
    r""" 统计空间频率 """
    dim = len(im.shape)
    if dim < 2 or dim > 3:
        raise "image is offen 2 or 3 dim, here is invalid"
    ans = []
    if dim == 2:
        v = spatial_freq_one_channel(im)
        ans.append(v)
    else:
        for i in range(im.shape[2]):
            v = spatial_freq_one_channel(im[:, :, i])
            ans.append(v)
    return ans


def info_entropy_one_channel(im):
    r""" 单分类计算信息熵 """
    # 创建一个eps,  0.000000001值,防止计算出现lg(0)
    eps = [1e-10] * 256
    # 不调用cv2方法,直接统计
    hist = np.array(eps)
    # 展开访问
    arr = im.reshape([-1])
    for v in arr:
        hist[v] = hist[v] + 1
    
    # 直方图/总像素数 获得概率 --- 具体概率怎么计算未知
    p = hist / (im.shape[0] * im.shape[1])
    # p * lg(p)
    p = p * np.log(p) * -1.0
    # - 求和
    H = np.sum(p)
    return H


def info_entropy(im):
    r""" 计算信息熵 """
    dim = len(im.shape)
    if dim < 2 or dim > 3:
        raise "image is offen 2 or 3 dim, here is invalid"
    ans = 0.0
    if dim == 2:
        v = info_entropy_one_channel(im)
        ans = ans + v
    else:
        for i in range(im.shape[2]):
            v = info_entropy_one_channel(im[:, :, i])
            ans = ans + v
    return ans
    


if __name__ == "__main__":
    file_name = "F:/480x320.yuv"
    out_dir = "F:/"
    h = 320  # 根据文件名中的宽高改写
    w = 480  # 根据文件名中的宽高改写
    nv21_patch_size = h * w * 3 // 2
    cnt = 0
    for patch in read_file_yile(file_name, nv21_patch_size):
        if patch:
            nv21 = load_nv21_image(patch, w, h)
            bgr = nv21_to_bgr(nv21)
            cv2.imwrite(out_dir + "/" + str(cnt) + ".jpg", bgr)
            cnt = cnt + 1
